#!/usr/bin/env python3
# coding=utf-8

import rospy
from geometry_msgs.msg import Twist
from geometry_msgs.msg import PoseStamped
from mavros_msgs.srv import *
from mavros_msgs.msg import State
import math
import sys, select, termios, tty
import cap
from geometry_msgs.msg import PointStamped

# 空格：降落
# 5  ：开启offboard模式
# 6  ：解锁，准备起飞
# 7  ：起飞

msg = """
Control Your Turtlebot!
---------------------------
Moving around:
   u    i    o
   j    k    l
   m    ,    .
   a    d   : move left/right
   r    f   : move up/down

q/z : increase/decrease max speeds by 10%
w/x : increase/decrease only linear speed by 10%
e/c : increase/decrease only angular speed by 10%
space key, k : force stop
anything else : stop smoothly
b : switch to OmniMode/CommonMode
CTRL-C to quit
"""
Omni = 0  # 全向移动模式

x = 0  # 前进后退方向
y = 0  # 左右移动方向
z = 0  # 上下移动方向
th = 0  # 转向/横向移动方向
count = 0  # 键值不再范围计数
target_speed = 0  # 前进后退目标速度
target_z_speed = 0  # 上下运动目标速度
target_turn = 0  # 转向目标速度
target_side_speed = 0  # 左右平移目标速度
control_speed = 0  # 前进后退实际控制速度
control_z_speed = 0  # 上下运动实际控制速度
control_turn = 0  # 转向实际控制速度
control_side_speed = 0  # 左右平移实际控制速度

# 键值对应移动/转向方向
moveBindings = {
    'i': (1, 0),  # 前进
    'o': (1, -1),  # 前右
    'j': (0, 1),  # 左转
    'l': (0, -1),  # 右转
    'u': (1, 1),  # 前左
    ',': (-1, 0),  # 后退
    '.': (-1, 1),  # 后右
    'm': (-1, -1),  # 后左
    'a': (0, -1),  # 左平移
    'd': (0, 1),  # 右平移
}

# 键值对应速度增量
speedBindings = {
    'q': (1.1, 1.1),
    'z': (0.9, 0.9),
    'w': (1.1, 1),
    'x': (0.9, 1),
    'e': (1, 1.1),
    'c': (1, 0.9),
}


def getKey():
    tty.setraw(sys.stdin.fileno())
    rlist, _, _ = select.select([sys.stdin], [], [], 0.1)
    if rlist:
        key = sys.stdin.read(1)
    else:
        key = ''

    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
    return key


speed = 0.2  # 默认移动速度 m/s
turn = 1  # 默认转向速度 rad/s


def vels(speed, turn):
    return "currently:\tspeed %s\tturn %s " % (speed, turn)


sita = 0.0  # 朝向
z = 0
w = 0
zf = 1
current_height = 0


def pose_cb(m):
    global sita
    global z
    global w
    global zf
    global current_height
    z = m.pose.orientation.z
    w = m.pose.orientation.w
    if z * w > 0:
        zf = 1
    else:
        zf = -1
    sita = 2 * math.acos(w) * 180 / math.pi
    current_height = m.pose.position.z



current_state = State()

armed_status = 0#解锁状态标志位

#接收惯导坐标
current_x=0
current_y=0
def position_callback(msg):
    # 打印接收到的坐标信息
    # rospy.loginfo(f"Received position: x={msg.point.x}, y={msg.point.y}, z={msg.point.z}")
    global current_x
    global current_y
    if msg is None or msg.point is None:
        print("Error: Received None message or point")
        return
    current_x = msg.point.x
    current_y = msg.point.y
def state_cb(state):
    global current_state
    global  armed_status
    current_state = state
    armed_status = current_state.armed

takeoff_flag=0
def takeoff(target_height):
    global z
    global current_height
    global takeoff_flag
    print("Preparing for takeoff...")

    setModeServer(custom_mode='AUTO.TAKEOFF')  # 设置模式为自动起飞
    print("Comfirmed takeoff")
    while current_height < target_height:
        z = 1  # 上升速度设置
        update_velocity()  # 更新速度
        rospy.sleep(0.1)  # 稍微休息一下，不要太心急
        if current_height >= target_height:
            print("Reached target height: {:.2f} meters. Leveling off.".format(current_height))
            z = 0  # 停止上升
            update_velocity()  # 最后一次速度更新，保证平稳
            break
    setModeServer(custom_mode='OFFBOARD')  # 设置模式为外部控制
    takeoff_flag=1

def update_velocity():
    twist = Twist()
    twist.linear.z = z
    pub.publish(twist)

def pid(kp,ki,kd,target,current,limit_speed):#速度限幅输入-1时不限幅
    if current is None:
        return 0  # 或者其他适当的默认值
    else:
        #定义相关变量
        error = 0
        diff = 0
        intrgrall = 0
        last_error=0
        output = 0
        #计算各控制项
        error = target-current
        diff = error-last_error
        intrgrall = intrgrall + error
        # 接近期望值时清除积分项
        if abs(error)<=0.01:
            intrgrall=0
        #计算总控制项
        output = kp*error+kd*diff+ki*intrgrall
        #输出速度限幅
        if limit_speed==-1 :
            return output
        else:
            if output>limit_speed:
                output = limit_speed
                return output
            elif output<-limit_speed:
                output = -limit_speed
                return  output


def height_keep(cruise_height,limit_speed):#输入巡航高度
    global current_height
    global z
    if cruise_height is not None and current_height is not None:
        z = pid(5, 2, 0, cruise_height, current_height, limit_speed)
    else:
        print("Error: Height data is missing")
        z = 0

def Waypoint_flight(waypoint_x,waypoint_y,height,limit_speed_x,limit_speed_y,limit_speed_z):
    global current_x
    global current_y
    global x   # 前进后退方向
    global y   # 左右移动方向
    global z   # 上下移动方向
    global th   # 转向/横向移动方向
    global count   # 键值不再范围计数
    global target_speed   # 前进后退目标速度
    global target_z_speed   # 上下运动目标速度
    global target_turn   # 转向目标速度
    global target_side_speed   # 左右平移目标速度
    global control_speed   # 前进后退实际控制速度
    global control_z_speed   # 上下运动实际控制速度
    global control_turn   # 转向实际控制速度
    global control_side_speed   # 左右平移实际控制速度

    print("go to")
    print(waypoint_x,waypoint_y,height)

    while (1):
        if current_x is not None and current_y is not None:
            x=pid(5,2,0,waypoint_x,current_x,limit_speed_x)#x方向航点飞行
            y=pid(5,2,0,waypoint_y,current_y,limit_speed_y)#y方向航点飞行
            height_keep(height,limit_speed_z)

            if x is None:
                x=0
            if y is None:
                y=0
            if z is None:
                z=0
            if th is None:
                th=0
            # rospy.loginfo(f"x={x}, y={y}, z={z}, th={th}")

            target_speed = speed * x
            target_z_speed = speed * z
            target_turn = turn * th
            target_side_speed = speed * y

            if target_speed > control_speed:
                control_speed = min(target_speed, control_speed + 0.1)
            elif target_speed < control_speed:
                control_speed = max(target_speed, control_speed - 0.1)
            else:
                control_speed = target_speed

            if target_z_speed > control_z_speed:
                control_z_speed = min(target_z_speed, control_z_speed + 0.1)
            elif target_z_speed < control_z_speed:
                control_z_speed = max(target_z_speed, control_z_speed - 0.1)
            else:
                control_z_speed = target_z_speed

            if target_turn > control_turn:
                control_turn = min(target_turn, control_turn + 0.5)
            elif target_turn < control_turn:
                control_turn = max(target_turn, control_turn - 0.5)
            else:
                control_turn = target_turn

            if target_side_speed > control_side_speed:
                control_side_speed = min(target_side_speed, control_side_speed + 0.1)
            elif target_side_speed < control_side_speed:
                control_side_speed = max(target_side_speed, control_side_speed - 0.1)
            else:
                control_side_speed = target_side_speed

            y_sita = math.sin(sita / 180 * math.pi)
            if y_sita < 0:
                y_sita = -y_sita
            y_sita = y_sita * zf

            twist = Twist()
            twist.linear.x = control_speed * math.cos(sita / 180 * math.pi)
            twist.linear.y = control_speed * y_sita + control_side_speed  # 包含左右平移
            twist.linear.z = control_z_speed
            twist.angular.x = 0
            twist.angular.y = 0
            twist.angular.z = control_turn

            pub.publish(twist)

            if abs(waypoint_x-current_x)<0.25 and abs(waypoint_y-current_y)<0.25 and abs(height-current_height)<0.25:#接近目标时结束循环
                break

        else:
            print("Error: Current position data is missing")

def Position_hold(waypoint_x,waypoint_y,height,limit_speed_x,limit_speed_y,keep_time):
    global current_x
    global current_y
    global x   # 前进后退方向
    global y   # 左右移动方向
    global z   # 上下移动方向
    global th   # 转向/横向移动方向
    global count   # 键值不再范围计数
    global target_speed   # 前进后退目标速度
    global target_z_speed   # 上下运动目标速度
    global target_turn   # 转向目标速度
    global target_side_speed   # 左右平移目标速度
    global control_speed   # 前进后退实际控制速度
    global control_z_speed   # 上下运动实际控制速度
    global control_turn   # 转向实际控制速度
    global control_side_speed   # 左右平移实际控制速度

    print("hold at")
    print(waypoint_x,waypoint_y,height,keep_time)

    start_time = rospy.Time.now()
    current_time = rospy.Time.now()

    while (1):
        current_time = rospy.Time.now()
        d_t=(current_time-start_time).to_sec()
        if current_x is not None and current_y is not None:
            x=pid(5,2,0,waypoint_x,current_x,limit_speed_x)#x方向航点飞行
            y=pid(5,2,0,waypoint_y,current_y,limit_speed_y)#y方向航点飞行
            height_keep(height,-1)

            if x is None:
                x=0
            if y is None:
                y=0
            if z is None:
                z=0
            if th is None:
                th=0
            # rospy.loginfo(f"x={x}, y={y}, z={z}, th={th}")

            target_speed = speed * x
            target_z_speed = speed * z
            target_turn = turn * th
            target_side_speed = speed * y

            if target_speed > control_speed:
                control_speed = min(target_speed, control_speed + 0.1)
            elif target_speed < control_speed:
                control_speed = max(target_speed, control_speed - 0.1)
            else:
                control_speed = target_speed

            if target_z_speed > control_z_speed:
                control_z_speed = min(target_z_speed, control_z_speed + 0.1)
            elif target_z_speed < control_z_speed:
                control_z_speed = max(target_z_speed, control_z_speed - 0.1)
            else:
                control_z_speed = target_z_speed

            if target_turn > control_turn:
                control_turn = min(target_turn, control_turn + 0.5)
            elif target_turn < control_turn:
                control_turn = max(target_turn, control_turn - 0.5)
            else:
                control_turn = target_turn

            if target_side_speed > control_side_speed:
                control_side_speed = min(target_side_speed, control_side_speed + 0.1)
            elif target_side_speed < control_side_speed:
                control_side_speed = max(target_side_speed, control_side_speed - 0.1)
            else:
                control_side_speed = target_side_speed

            y_sita = math.sin(sita / 180 * math.pi)
            if y_sita < 0:
                y_sita = -y_sita
            y_sita = y_sita * zf

            twist = Twist()
            twist.linear.x = control_speed * math.cos(sita / 180 * math.pi)
            twist.linear.y = control_speed * y_sita + control_side_speed  # 包含左右平移
            twist.linear.z = control_z_speed
            twist.angular.x = 0
            twist.angular.y = 0
            twist.angular.z = control_turn

            pub.publish(twist)

            if d_t>=keep_time :
                break

        else:
            print("Error: Current position data is missing")


task_flag=0
def flight_task():
    global  takeoff_flag
    global task_flag
    if takeoff_flag==0:
        takeoff(1)
    else:
        if task_flag == 0:
            Position_hold(0,0,1,2,2,2)#起飞后悬停2s稳定姿态

            print("approaching shelf 1")
            Waypoint_flight(0.5, 0, 1, 2, 2,2)#货架1进场
            Position_hold(0.5, 0, 1, 2, 2, 5)
            # 货架1扫描航路
            print("shelf 1 scan start")
            Waypoint_flight(0.5, 0.75, 1, 2, 2,2)
            Position_hold(0.5, 0.75, 1, 2, 2, 2)
            Waypoint_flight(0.5, 1.25, 1, 2, 2,2)
            Position_hold(0.5, 1.25, 1, 2, 2, 2)
            Waypoint_flight(0.5, 1.75, 1, 2, 2,2)
            Position_hold(0.5, 1.75, 1, 2, 2, 2)
            Waypoint_flight(0.5, 1.75, 1.4, 2, 2,2)
            Position_hold(0.5, 1.75, 1.4, 2, 2, 2)
            Waypoint_flight(0.5, 1.25, 1.4, 2, 2,2)
            Position_hold(0.5, 1.25, 1.4, 2, 2, 2)
            Waypoint_flight(0.5, 0.75, 1.4, 2, 2,2)
            Position_hold(0.5, 0.75, 1.4, 2, 2, 2)
            print("shelf 1 scan finish")

            #货架2转场航路
            print("leave shelf 1")
            Waypoint_flight(0.0, 0.75, 1.5, 2, 2,2)#货架1离场
            Position_hold(0.0, 0.75, 1.5, 2, 2, 1)
            Waypoint_flight(0.0, -0.3, 1.5, 2, 2,2)
            Position_hold(0.0, -0.3, 1.5, 2, 2, 1)
            print("go to shelf 2")
            Waypoint_flight(2.5, -0.3, 1.5, 2, 2,2)
            Position_hold(2.5, -0.3, 1.5, 2, 2, 1)
            

            #货架2扫描航路
            print("approaching shelf2")
            Waypoint_flight(2.5, -0.3, 1, 2, 2,2)#货架2进场
            Position_hold(2.5, -0.3, 1, 2, 2, 1)
            Waypoint_flight(2.5, 0.75, 1, 2, 2,2)
            print("shelf 2 scan start")
            Position_hold(2.5, 0.75, 1, 2, 2, 2)
            Waypoint_flight(2.5, 1.25, 1, 2, 2,2)
            Position_hold(2.5, 1.25, 1, 2, 2, 2)
            Waypoint_flight(2.5, 1.75, 1, 2, 2,2)
            Position_hold(2.5, 1.75, 1, 2, 2, 2)
            Waypoint_flight(2.5, 1.75, 1.4, 2, 2,2)
            Position_hold(2.5, 1.75, 1.4, 2, 2, 2)
            Waypoint_flight(2.5, 1.25, 1.4, 2, 2,2)
            Position_hold(2.5, 1.25, 1.4, 2, 2, 2)
            Waypoint_flight(2.5, 0.75, 1.4, 2, 2,2)
            Position_hold(2.5, 0.75, 1.4, 2, 2, 2)
            print("shelf 2 scan finish")

            #转场降落航路
            print("leave shelf 2")
            Waypoint_flight(2, 0.75, 1.4, 2, 2,2)#货架2离场
            Position_hold(2, 0.75, 1.4, 2, 2, 1)
            Waypoint_flight(2, 3.65, 1.4, 2, 2,2)
            Position_hold(2, 3.65, 1.4, 2, 2, 1)
            print("go to destination")#着陆场进场航路
            Waypoint_flight(3.5, 3.65, 1.4, 2, 2,2)
            Position_hold(3.5, 3.65, 1.4, 2, 2, 1)
            print("approaching tarmac")#着陆场进近航路
            Waypoint_flight(3.5, 2.9, 1.4, 2, 2,2)
            Position_hold(3.5, 2.9, 1.4, 2, 2, 2)
            Waypoint_flight(3.5, 2.9, 1, 2, 2,2)
            Position_hold(3.5, 2.9, 1, 2, 2, 2)
            Waypoint_flight(3.5, 2.9, 0.5, 2, 2,2)
            Position_hold(3.5, 2.9, 0.5, 2, 2, 8)
            #回原点，误差评估
            print("return to home")
            Waypoint_flight(0, 0, 0.5, 2, 2,2)
            Position_hold(0, 0, 0.5, 2, 2, 8)
            print("mission finish,landing now")
            setModeServer(custom_mode='AUTO.LAND')
            print("Vehicle Land")
            task_flag=1
        else:
            print("mission finish")


if __name__ == "__main__":
    settings = termios.tcgetattr(sys.stdin)

    rospy.init_node('turtlebot_teleop')
    pub = rospy.Publisher('mavros/setpoint_velocity/cmd_vel_unstamped', Twist, queue_size=5)
    # 订阅惯性导航坐标话题
    rospy.Subscriber('computed_position', PointStamped, position_callback)
    rospy.Subscriber('mavros/local_position/pose', PoseStamped, pose_cb)
    rospy.Subscriber('mavros/state', State, state_cb)

    setModeServer = rospy.ServiceProxy('mavros/set_mode', SetMode)
    armServer = rospy.ServiceProxy('/mavros/cmd/arming', CommandBool)


    try:
        print(msg)
        print(vels(speed, turn))
        armServer(True)
        print("Vehicle armed")
        # current_height=0#校正海压令当前高度为0
        # setModeServer(custom_mode='OFFBOARD')  # 设置模式为外部控制
        # takeoff(1.5)
        while (1):
            # print(current_height)
            key = getKey()
            if armed_status==1:
                flight_task()
                # break
            else:
                armServer(True)


            if key == ' ':
                print("Vehicle Land")
                setModeServer(custom_mode='AUTO.LAND')

            if (key == '\x03'):
                break



    except Exception as e:
        print(e)

    finally:
        twist = Twist()
        twist.linear.x = 0
        twist.linear.y = 0
        twist.linear.z = 0
        twist.angular.x = 0
        twist.angular.y = 0
        twist.angular.z = control_turn
        pub.publish(twist)

    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
